// Package ratelimit
/**
* @Project : geektime-basic-go-study
* @File    : builder.go
* @IDE     : GoLand
* @Author  : Tvux
* @Date    : 2024/9/21 18:10
**/

package ratelimit

import (
	_ "embed"
	"fmt"
	"github.com/HJH0924/GenericGo/ratelimiter"
	"github.com/gin-gonic/gin"
	"log"
	"net/http"
)

type Builder struct {
	prefix  string
	limiter ratelimiter.RateLimiter
}

func NewBuilder(limiter ratelimiter.RateLimiter) *Builder {
	return &Builder{
		prefix:  "ip-limiter",
		limiter: limiter,
	}
}

func (Self *Builder) Prefix(prefix string) *Builder {
	Self.prefix = prefix
	return Self
}

func (Self *Builder) limit(ctx *gin.Context) (bool, error) {
	key := fmt.Sprintf("%s:%s", Self.prefix, ctx.ClientIP())
	return Self.limiter.IsLimit(ctx, key)
}

func (Self *Builder) Build() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		limited, err := Self.limit(ctx)
		if err != nil {
			// 记录错误日志
			log.Printf("Error in rate limiting: %v", err)
			ctx.AbortWithStatus(http.StatusInternalServerError)
			return
		}
		if limited {
			// 记录限流日志，包括请求的 IP 地址和时间戳
			log.Printf("Rate limit exceeded for IP: %s", ctx.ClientIP())
			ctx.AbortWithStatus(http.StatusTooManyRequests)
			return
		}
		ctx.Next()
	}
}
